#include "mogic.h"
#include "console.h"
#include "sd.h"
#include "fs.h"
#include "file.h"
#include "initcode.h"
#include "spinlock.h"
#include "cop.h"
#include "log.h"

static struct {
    struct spinlock lock;
    int count;
    cop_task_t task[N_COWORKER];
} cop;

// for synchronization
static struct {
    int init;
    int begin;
    int end;
} _sync;

void cop_init() {
    initlock(&cop.lock, "cop.lock");

    cprintf("cop_init: _sync.init=%p\n", &_sync.init);
    cprintf("cop_init: _sync.begin=%p\n", &_sync.begin);
    cprintf("cop_init: _sync.end=%p\n", &_sync.end);
}

/**
 * wait for all coworkers
 */
static void wait_cop() {
    acquire(&cop.lock);
    while (cop.count < N_COWORKER)
        sleep(&_sync.init, &cop.lock);
    release(&cop.lock);
}

static int task_finished() {
    // acquire(&cop.lock);
    int ok = 1;
    for (int i = 0; i < N_COWORKER; i++) {
        ok &= cop.task[i] == 0;
    }
    // release(&cop.lock);

    return ok;
}

void task_dispatch(int n, cop_task_t task) {
    if (n < 0)
        n = N_COWORKER;

    acquire(&cop.lock);
    for (int i = 0; i < n && i < N_COWORKER; i++)
        cop.task[i] = task;
    wakeup(&_sync.begin);
    release(&cop.lock);
}

void task_wait() {
    acquire(&cop.lock);
    while (!task_finished())
        sleep(&_sync.end, &cop.lock);
    release(&cop.lock);
}

void task_run(int n, cop_task_t task) {
    if (n < 0)
        n = N_COWORKER + 1;

    if (n > 1)
        task_dispatch(n - 1, task);
    if (n > 0)
        task();

    task_wait();
}

static void say_happy_new_year() {
    cprintf("HAPPY NEW YEAR！CPU %d 给您拜年了！\n", arm_cpuid());
}

int sys_cop() {
    acquire(&cop.lock);
    int idx = cop.count++;
    wakeup(&_sync.init);
    release(&cop.lock);

    cprintf("cpu %d: cop: coworker #%d started.\n", arm_cpuid(), idx);

    while (1) {
        acquire(&cop.lock);
        while (!cop.task[idx])
            sleep(&_sync.begin, &cop.lock);
        release(&cop.lock);

        cop.task[idx]();

        acquire(&cop.lock);
        cop.task[idx] = 0;
        wakeup(&_sync.end);
        release(&cop.lock);
    }

    return 0;
}

/**
 * run *init() that requires sleep & wakeup.
 */
int sys_mogic() {
    static int life = 2, tea = 0;
    life--;
    tea++;

    if (life > 0) {
        wait_cop();
        cprintf("cpu %d: mogic: the Mogician wakes up!\n", arm_cpuid());

        task_run(3, say_happy_new_year);

        // cprintf("sys_sdtest: SD card test started.\n");
        // task_run(3, sd_test);

        fsinit();

#ifdef __MORE_TEST
        log_test();
        fs_test();
#endif

        task_dispatch(-1, exit);

        for (int i = 0; i < N_COWORKER; i++) {
            assert(wait() >= 0);
        }
    }

    return 0;
}
